package com.haotian.coder.ai;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;

import com.haotian.coder.AiCodeGeneratorService;
import com.haotian.coder.guardrail.PromptSafetyInputGuardrail;
import com.haotian.coder.model.enums.CodeGenTypeEnum;
import com.haotian.coder.service.ChatHistoryService;
import com.haotian.coder.tools.ToolManager;
import dev.langchain4j.community.store.memory.chat.redis.RedisChatMemoryStore;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.service.AiServices;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.concurrent.TimeUnit;

@Slf4j
@Configuration
public class AiCodeGeneratorServiceFactory {

    /**
     * 缓存策略
     * 写入后30分钟过期
     * 读取后10分钟过期
     */
    private final Cache<String, AiCodeGeneratorService> serviceCache = Caffeine.newBuilder().maximumSize(1000).expireAfterWrite(30, TimeUnit.MINUTES).expireAfterAccess(10, TimeUnit.MINUTES).removalListener((key, value, cause) -> {
        log.info("缓存: {} 被删除，原因{}", key, cause);
    }).build();

    @Resource
    ChatHistoryService chatHistoryService;

    @Resource
    private ChatModel chatModel;

    @Resource
    @Qualifier("openAiStreamingChatModel")
    private StreamingChatModel openAiStreamChatModel;

    @Resource
    @Qualifier("reasoningStreamingChatModel")
    private StreamingChatModel reasoningStreamingChatMode;

    @Resource
    private RedisChatMemoryStore redisChatMemoryStore;

    @Resource
    private ToolManager toolManager;

    /**
     * 默认提供一个 Bean
     */
    @Bean
    public AiCodeGeneratorService aiCodeGeneratorService() {
        return getAiCodeGeneratorService(0L);
    }

    /**
     * 根据 appId 获取服务（带缓存）这个方法是为了兼容历史逻辑
     */
    public AiCodeGeneratorService getAiCodeGeneratorService(long appId) {
        return getAiCodeGeneratorService(appId, CodeGenTypeEnum.HTML);
    }


    /**
     * 根据appID创建ai服务,缓存模式
     *
     * @param appID
     * @return
     */
    public AiCodeGeneratorService getAiCodeGeneratorService(Long appID, CodeGenTypeEnum codeGenTypeEnum) {
        String cacheKey = buildCacheKey(appID, codeGenTypeEnum);
        return serviceCache.get(cacheKey, func -> createAiCodeGeneratorService(appID, codeGenTypeEnum));
    }

    private String buildCacheKey(Long appId, CodeGenTypeEnum codeGenTypeEnum) {
        return appId + "_" + codeGenTypeEnum.getValue();
    }

    public AiCodeGeneratorService createAiCodeGeneratorService(Long appID, CodeGenTypeEnum codeGenTypeEnum) {

        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                .id(appID)
                .maxMessages(100)
                .chatMemoryStore(redisChatMemoryStore)
                .build();

        chatHistoryService.loadHistoryMemory(appID, chatMemory, 50);


        return switch (codeGenTypeEnum) {

            case VUE_PROJECT ->
                    AiServices.builder(AiCodeGeneratorService.class)
                            .streamingChatModel(reasoningStreamingChatMode)
                            .tools(toolManager.getAllTools())
                            .chatMemoryProvider(memoryId -> chatMemory)
                            .inputGuardrails(new PromptSafetyInputGuardrail()) // 护轨，审查prompt 是否安全
                            .hallucinatedToolNameStrategy(toolExecutionRequest -> ToolExecutionResultMessage.from(toolExecutionRequest, "Error: there is no tool called " + toolExecutionRequest.name()))
                            .build();

            case HTML, MULTI_FILE ->
                    AiServices.builder(AiCodeGeneratorService.class)
                            .chatModel(chatModel)
                            .streamingChatModel(openAiStreamChatModel)
                            .inputGuardrails(new PromptSafetyInputGuardrail())
                            .chatMemory(chatMemory)
                            .inputGuardrails(new PromptSafetyInputGuardrail()) // 护轨，审查prompt 是否安全
                            .build();
        };

    }

}

